"""
model name : 目标跟踪
file       : 1_OTB数据标签转YOLO标签.py
information:
    author : OuYang
    time   : 2025/7/19
"""
import os
import cv2

input_dir = "D:/Downloads/Car24"
output_dir = "d:/Downloads/Car24_Yolo"


def convert_otb_to_yolo(input_dir, output_dir):
    label_file = os.path.join(input_dir, "groundtruth_rect.txt")
    image_dir = os.path.join(input_dir, "img")

    if not os.path.exists(label_file) or not os.path.exists(image_dir):
        raise Exception("Not found label or image dir")

    # 获取标签内的每一行
    with open(label_file, "r") as file:
        lines = file.readlines()

    for i, line in enumerate(lines):
        # 获取每一帧图片
        image_path = os.path.join(image_dir, f"{i + 1:04d}.jpg")
        print(image_path)

        if not os.path.exists(image_path):
            raise Exception("image_path does not exist")

        # 获取每一帧的标签数据
        values = list(map(float, line.strip().split(',')))
        x, y, w, h = values[:4]

        # 获取图片大小
        image = cv2.imread(image_path)
        img_h, img_w = image.shape[:2]

        # 转换伟YOLO格式
        x_center = (x + w / 2) / img_w
        y_center = (y + h / 2) / img_h
        w_norm = w / img_w
        h_norm = h / img_h

        yolo_label = os.path.join(output_dir, 'label', f"{i + 1:04d}.txt")
        os.makedirs(os.path.dirname(yolo_label), exist_ok=True)
        with open(yolo_label, "w") as file:
            file.write(f'0 {x_center: .6f} {y_center: .6f} {w_norm: .6f} {h_norm: .6f}')


if __name__ == '__main__':
    convert_otb_to_yolo(input_dir, output_dir)